from pprint import pprint
import pandas as pd
import pymysql


class DataOperation:
    def __init__(self, file, path=None, sql=None, cn=None):
        if file == 'csv':
            self.data = pd.read_csv(path)
        elif file == 'excel':
            self.data = pd.read_excel(path, sheet_name=1)
        elif file == 'mysql':
            self.data = pd.read_sql(sql, cn)
        else:
            self.data = None

    # 转化为字典
    def get_data_todict(self):
        return [self.data.loc[i].to_dict() for i in self.data.index.values]

    # 转化为列表
    def get_data_to_list(self):
        return self.data.values.tolist()


if __name__ == '__main__':
    # do = DataOperation('csv', '../data/lj_data.csv')
    # pprint(do.get_data_todict())
    # pprint(do.get_data_to_list())

    cn = pymysql.connect(
        host='localhost',
        user='root',
        password='kaikai',
        database='news',
        port=3306,
        charset='utf8'
    )
    sql = 'select * from company_news'
    do = DataOperation('mysql', sql=sql, cn=cn, path=None)
    print(do.data.head(2))


